{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71a329f0",
   "metadata": {},
   "source": [
    "# PySDK instruction for using LMI container on SageMaker\n",
    "In this tutorial, you will use LMI container from DLC to SageMaker and run inference with it.\n",
    "\n",
    "Please make sure the following permission granted before running the notebook:\n",
    "\n",
    "- S3 bucket push access\n",
    "- SageMaker access\n",
    "\n",
    "## Step 1: Let's bump up SageMaker and import stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67fa3208",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install sagemaker boto3 awscli --upgrade  --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9ac353",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker import image_uris, serializers, deserializers\n",
    "from sagemaker.djl_inference.model import DeepSpeedModel, DJLModel, HuggingFaceAccelerateModel\n",
    "\n",
    "role = sagemaker.get_execution_role()  # execution role for the endpoint\n",
    "sess = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs\n",
    "region = sess._region_name  # region name of the current SageMaker Studio environment\n",
    "account_id = sess.account_id()  # account_id of the current SageMaker Studio environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53bb86c4",
   "metadata": {},
   "source": [
    "### (Remove if not needed) upload HuggingFace model to S3 bucket\n",
    "\n",
    "LMI has good capability to download model in a S3 bucket. This step is recommended if you would like to speed up model loading process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c446e156",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install huggingface_hub --upgrade --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c644d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "# - This will download the model into the ./model directory where ever the jupyter file is running\n",
    "local_model_path = Path(\"/tmp\")\n",
    "local_model_path.mkdir(exist_ok=True)\n",
    "model_name = \"facebook/opt-6.7b\"\n",
    "# Only download pytorch checkpoint files\n",
    "allow_patterns = [\"*.json\", \"*.pt\", \"*.bin\", \"*.txt\", \"*.model\"]\n",
    "\n",
    "# - Leverage the snapshot library to donload the model since the model is stored in repository using LFS\n",
    "model_download_path = snapshot_download(\n",
    "    repo_id=model_name,\n",
    "    cache_dir=local_model_path,\n",
    "    allow_patterns=allow_patterns,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baf63383",
   "metadata": {},
   "source": [
    "### (remove if not needed) DeepSpeed HF faster loading\n",
    "\n",
    "DeepSpeed offers a way to speed up model loading while keep the CPU memory low. This has only been tested with\n",
    "- OPT\n",
    "- GPT-NeoX\n",
    "- BLOOM\n",
    "\n",
    "We just put a checkpoints.json along with it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "700afc04",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import json\n",
    "checkpoints_json = os.path.join(model_download_path, \"checkpoints.json\")\n",
    "tensor_parallel_degree=4\n",
    "weight_dtype=\"float16\"\n",
    "\n",
    "with io.open(checkpoints_json, \"w\", encoding=\"utf-8\") as f:\n",
    "    file_list = [str(entry).split('/')[-1] for entry in Path(model_download_path).rglob(\"*.[bp][it][n]\") if entry.is_file()]\n",
    "    data = {\"type\": \"ds_model\", \"checkpoints\": file_list, \"version\": 1.0, \"parallelization\": \"tp\", \"tp_size\": tensor_parallel_degree, \"dtype\": weight_dtype}\n",
    "    json.dump(data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5105410e",
   "metadata": {},
   "source": [
    "Upload the model to S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a6210dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "key_prefix=\"lmi_models/mymodel\"\n",
    "model_artifact = sess.upload_data(path=model_download_path, key_prefix=key_prefix)\n",
    "print(f\"Model uploaded to --- > {model_artifact}\")\n",
    "print(f\"You can set option.s3url={model_artifact}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e58cf33",
   "metadata": {},
   "source": [
    "## Step 2: Start building SageMaker endpoint\n",
    "In this step, we will build SageMaker endpoint from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d955679",
   "metadata": {},
   "source": [
    "### Getting the container image URI\n",
    "\n",
    "Available framework are:\n",
    "- djl-deepspeed (0.20.0, 0.21.0)\n",
    "- djl-fastertransformer (0.21.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a174b36",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_uri = image_uris.retrieve(\n",
    "        framework=\"djl-deepspeed\",\n",
    "        region=sess.boto_session.region_name,\n",
    "        version=\"0.21.0\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11601839",
   "metadata": {},
   "source": [
    "### Create SageMaker model\n",
    "\n",
    "Here we are using [LMI PySDK](https://sagemaker.readthedocs.io/en/stable/frameworks/djl/using_djl.html) to create the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b1e5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not model_artifact:\n",
    "    model_artifact = \"s3://my_bucket/my_saved_model_artifacts/\"\n",
    "model = DeepSpeedModel(\n",
    "    model_artifact,\n",
    "    role,\n",
    "    data_type=\"fp16\",\n",
    "    task=\"text-generation\",\n",
    "    tensor_parallel_degree=2, # number of gpus to partition the model across using tensor parallelism\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "004f39f6",
   "metadata": {},
   "source": [
    "### Create SageMaker endpoint\n",
    "\n",
    "You need to specify the instance to use and endpoint names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0e61cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_type = \"ml.g4dn.4xlarge\"\n",
    "endpoint_name = sagemaker.utils.name_from_base(\"lmi-model\")\n",
    "\n",
    "model.deploy(initial_instance_count=1,\n",
    "             instance_type=instance_type,\n",
    "             endpoint_name=endpoint_name,\n",
    "             # container_startup_health_check_timeout=3600\n",
    "            )\n",
    "\n",
    "# our requests and responses will be in json format so we specify the serializer and the deserializer\n",
    "predictor = sagemaker.Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sess,\n",
    "    serializer=serializers.JSONSerializer(),\n",
    "    deserializer=deserializers.JSONDeserializer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb63ee65",
   "metadata": {},
   "source": [
    "## Step 3: Test and benchmark the inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bcef095",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n3 -r1\n",
    "predictor.predict(\n",
    "    {\"inputs\": \"Large model inference is\", \"parameters\": {}}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1cd9042",
   "metadata": {},
   "source": [
    "## Clean up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d674b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.delete_endpoint(endpoint_name)\n",
    "sess.delete_endpoint_config(endpoint_name)\n",
    "model.delete_model()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
